#include <sys/wait.h>
#include <iostream>
#include <unistd.h>
#include <string>
#include <cstdlib>
#include <vector>
#include <ctime>
#include "task.hpp"
using namespace std;
enum 
{
    UsageError = 1,
    ArgError,
    PipeError
};

// 站在父进程视角，抽象管道
class Channel
{
public:
    Channel(int wfd, pid_t sub_prcess_id, const string& name)
        :_wfd(wfd), _sub_process_id(sub_prcess_id), _name(name)
    {}
    
    void PrintDebug()
    {
        cout << "_wfd: " << _wfd;
        cout << ", _sub_process_id: " << _sub_process_id;
        cout << ", _name: " << _name << endl;
    }
    const string& name()
    {
        return _name;
    }
    int wfd()
    {
        return _wfd;
    }
    
    pid_t pid()
    {
        return _sub_process_id;
    }
    
    void Close()
    {
        close(_wfd);
    }
    ~Channel()
    {}
private:
    int _wfd; // 该管道的写端
    pid_t _sub_process_id; // 该管道对应的子进程
    string _name; // 管道名字 
};

// 提示用法：执行的进程名 + 需要创建的子进程数量
void Usage(const string& proc)
{
    cout << "Usage: " << proc << "subProcessNum" << endl;
}

// 对进程池进行抽象
class ProcessPool
{
public:
    ProcessPool(int sub_process_num)
        :_sub_process_num(sub_process_num)
    {}

    // 创建子进程和管道
    int create_prcess(work_t work)
    {
        vector<int> fds;
        for(int num = 0; num < _sub_process_num; num++)
        {
            int pipefd[2] = { 0 };
            int n = pipe(pipefd);
            if(n < 0) return PipeError;
    
            // 创建子进程
            pid_t id = fork();
            if(id == 0)
            {
                // 子进程 -> read 
                close(pipefd[1]); // 关闭对应管带的写端
                
                // 关闭从父进程继承下来上次的写端
                if(!fds.empty())
                {
                    for(auto fd: fds)
                    {
                        close(fd);
                    }
                }
                // 重定向，将键盘的内容输入到管道中
                dup2(pipefd[0], 0);
                work(pipefd[0]); // 向每一个子进程中传入其对应的管道读端，以方便得到具体任务是从哪得到的指令并由谁完成
                exit(0);
            }
            
            string cname = "channel-" + to_string(num);
    
            // 父进程 -> write
            close(pipefd[0]); // 关闭读端
    
            _channels.push_back(Channel(pipefd[1], id, cname));
            fds.push_back(pipefd[1]);
        }
        return 0;
    }
    // 选择一个管道
    int NextChannel()
    {
        static int next = 0;
        int c = next++; // 移动到下一个管道
        next %= _channels.size();
        return c;
    }
    
    // 父进程向管道写入一个任务下标
    void SendTaskCode(int index, uint32_t code)
    {
        // 打印信息 
        cout << "send code: " << code << " to "  << _channels[index].name()  << " sub process id: " << _channels[index].pid()  << endl;
        write(_channels[index].wfd(), &code, sizeof(code));
    }

    // 打印管道信息
    void Degub()
    {
        for(auto& e : _channels)
        {
            e.PrintDebug();
        }
    }
    // 将该进程池中所有的子进程退出任务
    void KillAll()
    {
        // 当一个管带的写端被关闭时，该管道对应的子进程就会被关闭
        // 遍历管道，关闭对应的写端
        for(auto& channel : _channels)
        {
            channel.Close();
            cout << "channel_wfd-" << channel.wfd() << " is closed and process of " << channel.pid() << " is about to exit --> " << endl;
            // 当管道被关闭后，就直接回收子进程资源
            pid_t id = channel.pid();
            pid_t rid = waitpid(id, nullptr, 0);
            if(rid == id)
            {
                cout << "The resources of process whose id is " << id << " are reclaimed" << endl;
            }
        }

    }
    
    // 父进程回收进程池中所有子进程的资源
    void Wait()
    {
//        for(auto& channel : _channels)
//        {
//            pid_t id = channel.pid();
//            pid_t rid = waitpid(id, nullptr, 0);
//            if(rid == id)
//            {
//                cout << "The resources of process whose id is " << id << " are reclaimed" << endl;
//            }
//        }
    }
private:
    int _sub_process_num; // 子进程数量
    vector<Channel> _channels; // 组织管道
};

// 控制进程池：指定一个进程池，完成cnt次任务
void CtrlProcessPool(ProcessPool* processpool_ptr, int cnt)
{
    while(cnt)
    {
        // 1. 选择一个管道
        int channel = processpool_ptr->NextChannel();

        // 2. 选择一个任务
        uint32_t code = NextTask();
        
        // 3. 发送任务
        processpool_ptr->SendTaskCode(channel, code);
        cnt--;
        sleep(1);
    }
}
int main(int argc, char* argv[])
{
    if(argc != 2)
    {
        Usage(argv[0]);
        return UsageError;
    }
    int sub_process_num = std::stoi(argv[1]);
    // 如果创建子进程小于等于0个，直接退出
    if(sub_process_num <= 0)
        return ArgError;
    srand((uint64_t)time(nullptr)); 
    // 1. 创建管道和子进程
    ProcessPool* processpool_ptr = new ProcessPool(sub_process_num);
    processpool_ptr->create_prcess(worker);
    processpool_ptr->Degub();
    // 2. 控制子进程
    CtrlProcessPool(processpool_ptr, 6);
    cout << "All task completed !!" << endl;
    // 3.回收子进程
    // 如果父进程没有回收子进程，父进程退出后，所有子进程会被init进程收养并关闭
    // a. 先将所有子进程退出
    processpool_ptr->KillAll();
    // b. 使用waitpid回收资源
    processpool_ptr->Wait();
    delete processpool_ptr;
    return 0;
}
